// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/check_duplicate_keys.h>

namespace kumo::pollux::functions {
    namespace {
        static const char *kNullKeyErrorMessage = "map key cannot be null";
        static const char *kIndeterminateKeyErrorMessage =
                "map key cannot be indeterminate";

        template<bool AllowDuplicateKeys>
        class MapFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                // No arguments case (empty map).
                if (args.empty()) {
                    auto emptyMapVector = std::make_shared<MapVector>(
                        context.pool(),
                        outputType,
                        nullptr, // nulls
                        rows.end(),
                        allocateOffsets(rows.end(), context.pool()),
                        allocateSizes(rows.end(), context.pool()),
                        BaseVector::create(outputType->childAt(0), 0, context.pool()),
                        BaseVector::create(outputType->childAt(1), 0, context.pool()));

                    context.moveOrCopyResult(emptyMapVector, rows, result);
                    return;
                }

                POLLUX_CHECK_EQ(args.size(), 2);

                auto keys = args[0];
                auto values = args[1];

                exec::DecodedArgs decodedArgs(rows, args, context);
                auto decodedKeys = decodedArgs.at(0);
                auto decodedValues = decodedArgs.at(1);

                static const char *kArrayLengthsMismatch =
                        "Key and value arrays must be the same length";
                auto checkNullsInKey =
                        [&](const auto &keysElements, auto offset, auto size) {
                    for (auto i = 0; i < size; ++i) {
                        POLLUX_USER_CHECK(
                            !keysElements->is_null_at(offset + i), kNullKeyErrorMessage);

                        POLLUX_USER_CHECK(
                            !keysElements->contains_null_at(offset + i),
                            "{}: {}",
                            kIndeterminateKeyErrorMessage,
                            keysElements->toString(offset + i));
                    }
                };
                // When context.throwOnError is false, some rows will be marked as
                // 'failed'. These rows should not be processed further. 'remainingRows'
                // will contain a subset of 'rows' that have passed all the checks (e.g.
                // keys are not nulls and number of keys and values is the same).
                exec::LocalSelectivityVector remainingRows(context, rows);

                // If both vectors have identity mapping, check if we can take the zero-copy
                // fast-path.
                if (decodedKeys->isIdentityMapping() &&
                    decodedValues->isIdentityMapping() &&
                    canTakeFastPath(
                        keys->as<ArrayVector>(), values->as<ArrayVector>(), rows)) {
                    auto keysArray = keys->as<ArrayVector>();
                    auto valuesArray = values->as<ArrayVector>();

                    // Verify there are no null keys.
                    auto keysElements = keysArray->elements();
                    if (keysElements->may_have_nulls() ||
                        keysElements->may_have_nulls_recursive()) {
                        context.applyToSelectedNoThrow(rows, [&](auto row) {
                            auto offset = keysArray->offsetAt(row);
                            auto size = keysArray->sizeAt(row);
                            checkNullsInKey(keysElements, offset, size);
                        });
                    }

                    context.deselectErrors(*remainingRows);

                    auto map_vector = std::make_shared<MapVector>(
                        context.pool(),
                        outputType,
                        keysArray->nulls(),
                        rows.end(),
                        keysArray->offsets(),
                        keysArray->sizes(),
                        keysArray->elements(),
                        valuesArray->elements());

                    if constexpr (!AllowDuplicateKeys) {
                        checkDuplicateKeys(map_vector, *remainingRows, context);
                    }
                    context.moveOrCopyResult(map_vector, rows, result);
                } else if (decodedKeys->isConstantMapping()) {
                    // Constant keys.
                    auto keysIndex = decodedKeys->index(rows.begin());
                    auto valueIndices = decodedValues->indices();

                    auto keysArray = decodedKeys->base()->as<ArrayVector>();
                    auto valuesArray = decodedValues->base()->as<ArrayVector>();

                    // Verify there are no null keys and no duplicate keys.
                    auto numKeys = keysArray->sizeAt(keysIndex);
                    auto keysElements = keysArray->elements();
                    auto keysOffset = keysArray->offsetAt(keysIndex);

                    // Sort indices of keys so that values at the indices are in ascending
                    // order. Then compare adjacent values through these indices to check for
                    // duplicate keys.
                    std::vector<vector_size_t> sortedIndices(numKeys);
                    std::iota(sortedIndices.begin(), sortedIndices.end(), keysOffset);
                    keysElements->sort_indices(sortedIndices, CompareFlags());
                    try {
                        if (keysElements->may_have_nulls() ||
                            keysElements->may_have_nulls_recursive()) {
                            checkNullsInKey(keysElements, keysOffset, numKeys);
                        }

                        if constexpr (!AllowDuplicateKeys) {
                            checkDuplicateConstantKeys(sortedIndices, keysElements);
                        }
                    } catch (const std::exception &) {
                        context.setErrors(rows, std::current_exception());
                    }

                    // Check array lengths
                    context.applyToSelectedNoThrow(*remainingRows, [&](vector_size_t row) {
                        POLLUX_USER_CHECK_EQ(
                            numKeys,
                            valuesArray->sizeAt(valueIndices[row]),
                            "{}",
                            kArrayLengthsMismatch);
                    });

                    context.deselectErrors(*remainingRows);

                    vector_size_t totalElements = remainingRows->countSelected() * numKeys;

                    BufferPtr offsets = allocateOffsets(rows.end(), context.pool());
                    auto rawOffsets = offsets->asMutable<vector_size_t>();

                    BufferPtr sizes = allocateSizes(rows.end(), context.pool());
                    auto rawSizes = sizes->asMutable<vector_size_t>();

                    BufferPtr keysIndices = allocate_indices(totalElements, context.pool());
                    auto rawKeysIndices = keysIndices->asMutable<vector_size_t>();

                    BufferPtr valuesIndices = allocate_indices(totalElements, context.pool());
                    auto rawValuesIndices = valuesIndices->asMutable<vector_size_t>();

                    vector_size_t offset = 0;
                    remainingRows->applyToSelected([&](vector_size_t row) {
                        rawOffsets[row] = offset;
                        rawSizes[row] = numKeys;

                        auto valuesOffset = valuesArray->offsetAt(valueIndices[row]);
                        for (vector_size_t i = 0; i < numKeys; i++) {
                            // Make keys in the result vector sorted to optimize subsequent
                            // processing on this result.
                            rawKeysIndices[offset + i] = sortedIndices[i];
                            rawValuesIndices[offset + i] =
                                    valuesOffset + (sortedIndices[i] - keysOffset);
                        }

                        offset += numKeys;
                    });

                    auto wrappedKeys = BaseVector::wrap_in_dictionary(
                        nullptr, keysIndices, totalElements, keysArray->elements());

                    auto wrappedValues = BaseVector::wrap_in_dictionary(
                        nullptr, valuesIndices, totalElements, valuesArray->elements());

                    auto map_vector = std::make_shared<MapVector>(
                        context.pool(),
                        outputType,
                        nullptr,
                        rows.end(),
                        offsets,
                        sizes,
                        wrappedKeys,
                        wrappedValues,
                        std::nullopt,
                        true);
                    context.moveOrCopyResult(map_vector, *remainingRows, result);
                } else {
                    auto keyIndices = decodedKeys->indices();
                    auto valueIndices = decodedValues->indices();

                    auto keysArray = decodedKeys->base()->as<ArrayVector>();
                    auto valuesArray = decodedValues->base()->as<ArrayVector>();

                    // Verify there are no null keys.
                    auto keysElements = keysArray->elements();
                    if (keysElements->may_have_nulls() ||
                        keysElements->may_have_nulls_recursive()) {
                        context.applyToSelectedNoThrow(*remainingRows, [&](auto row) {
                            auto offset = keysArray->offsetAt(keyIndices[row]);
                            auto size = keysArray->sizeAt(keyIndices[row]);
                            checkNullsInKey(keysElements, offset, size);
                        });
                        context.deselectErrors(*remainingRows);
                    }

                    // Check array lengths
                    context.applyToSelectedNoThrow(*remainingRows, [&](vector_size_t row) {
                        POLLUX_USER_CHECK_EQ(
                            keysArray->sizeAt(keyIndices[row]),
                            valuesArray->sizeAt(valueIndices[row]),
                            "{}",
                            kArrayLengthsMismatch);
                    });

                    context.deselectErrors(*remainingRows);

                    vector_size_t totalElements = 0;
                    remainingRows->applyToSelected([&](auto row) {
                        totalElements += keysArray->sizeAt(keyIndices[row]);
                    });

                    BufferPtr offsets = allocateOffsets(rows.end(), context.pool());
                    auto rawOffsets = offsets->asMutable<vector_size_t>();

                    BufferPtr sizes = allocateSizes(rows.end(), context.pool());
                    auto rawSizes = sizes->asMutable<vector_size_t>();

                    BufferPtr valuesIndices = allocate_indices(totalElements, context.pool());
                    auto rawValuesIndices = valuesIndices->asMutable<vector_size_t>();

                    BufferPtr keysIndices = allocate_indices(totalElements, context.pool());
                    auto rawKeysIndices = keysIndices->asMutable<vector_size_t>();

                    vector_size_t offset = 0;
                    remainingRows->applyToSelected([&](vector_size_t row) {
                        auto size = keysArray->sizeAt(keyIndices[row]);
                        rawOffsets[row] = offset;
                        rawSizes[row] = size;

                        auto keysOffset = keysArray->offsetAt(keyIndices[row]);
                        auto valuesOffset = valuesArray->offsetAt(valueIndices[row]);
                        for (vector_size_t i = 0; i < size; i++) {
                            rawKeysIndices[offset + i] = keysOffset + i;
                            rawValuesIndices[offset + i] = valuesOffset + i;
                        }

                        offset += size;
                    });

                    auto wrappedKeys = BaseVector::wrap_in_dictionary(
                        nullptr, keysIndices, totalElements, keysArray->elements());

                    auto wrappedValues = BaseVector::wrap_in_dictionary(
                        nullptr, valuesIndices, totalElements, valuesArray->elements());

                    auto map_vector = std::make_shared<MapVector>(
                        context.pool(),
                        outputType,
                        nullptr,
                        rows.end(),
                        offsets,
                        sizes,
                        wrappedKeys,
                        wrappedValues);
                    if constexpr (!AllowDuplicateKeys) {
                        checkDuplicateKeys(map_vector, *remainingRows, context);
                    }
                    context.moveOrCopyResult(map_vector, *remainingRows, result);
                }
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // array(K), array(V) -> map(K,V)
                // () -> map()
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("K")
                    .typeVariable("V")
                    .returnType("map(K,V)")
                    .argumentType("array(K)")
                    .argumentType("array(V)")
                    .build(),
                    exec::FunctionSignatureBuilder()
                    .returnType("map(unknown,unknown)")
                    .build(),
                };
            }

        private:
            // Can only take the fast path if keys and values have an equal
            // number of arrays and the offsets and sizes of these arrays match
            // 1:1. The map must be well formed for all elements, also ones not
            // in 'rows' in apply(). This is because canonicalize() will touch
            // all elements in any case.
            bool canTakeFastPath(
                ArrayVector *keys,
                ArrayVector *values,
                const SelectivityVector &rows) const {
                POLLUX_CHECK_GE(keys->size(), rows.end());
                POLLUX_CHECK_GE(values->size(), rows.end());
                // the fast path takes a reference to the keys and values and the
                // offsets and sizes from keys. This is valid only if the keys and
                // values align for all rows for both size and offset. Anything
                // else will break canonicalize().
                if (keys->size() != values->size()) {
                    return false;
                }
                for (auto row = 0; row < keys->size(); ++row) {
                    if (keys->is_null_at(row)) {
                        continue;
                    }

                    if (values->is_null_at(row)) {
                        return false;
                    }

                    if (keys->offsetAt(row) != values->offsetAt(row) ||
                        keys->sizeAt(row) != values->sizeAt(row)) {
                        return false;
                    }
                }
                return true;
            }

            std::optional<vector_size_t> findDuplicateKeys(
                const std::vector<vector_size_t> &sortedIndices,
                const VectorPtr &keysElements) const {
                for (auto i = 1; i < sortedIndices.size(); ++i) {
                    if (keysElements->equal_value_at(
                        keysElements.get(), sortedIndices[i], sortedIndices[i - 1])) {
                        return sortedIndices[i];
                    }
                }
                return std::nullopt;
            }

            void checkDuplicateConstantKeys(
                const std::vector<vector_size_t> &sortedIndices,
                const VectorPtr &keysElements) const {
                static const char *kDuplicateKey =
                        "Duplicate map keys ({}) are not allowed";

                if (auto duplicateIndex = findDuplicateKeys(sortedIndices, keysElements)) {
                    auto duplicateKey = keysElements->wrapped_vector()->toString(
                        keysElements->wrapped_index(duplicateIndex.value()));
                    POLLUX_USER_FAIL(kDuplicateKey, duplicateKey);
                }
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_map,
        MapFunction</*AllowDuplicateKeys=*/false>::signatures(),
        std::make_unique<MapFunction</*AllowDuplicateKeys=*/false>>());

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_map_allow_duplicates,
        MapFunction</*AllowDuplicateKeys=*/true>::signatures(),
        std::make_unique<MapFunction</*AllowDuplicateKeys=*/true>>());

    void registerMapFunction(const std::string &name, bool allowDuplicateKeys) {
        if (allowDuplicateKeys) {
            POLLUX_REGISTER_VECTOR_FUNCTION(udf_map_allow_duplicates, name);
        } else {
            POLLUX_REGISTER_VECTOR_FUNCTION(udf_map, name);
        }
    }
} // namespace kumo::pollux::functions
